from flask import Flask, Response, jsonify, request, render_template
import cv2
import numpy as np
import logging
from flask_cors import CORS

app = Flask(__name__)
CORS(app)  # 允许所有域名的请求

logging.basicConfig(level=logging.INFO)

latest_frame = None
STREAM_URL = "rtmp://120.46.37.219:3030/live/2"

@app.route('/')
def index():
    return render_template('index.html')

def generate_frames(stream_url):
    cap = cv2.VideoCapture(stream_url)
    if not cap.isOpened():
        app.logger.error(f"Cannot open stream: {stream_url}")
        return

    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 480)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    frame_skip = 5  # 跳过的帧数
    frame_count = 0

    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            app.logger.error("Failed to read frame from stream")
            break

        if frame_count % frame_skip == 0:
            ret, buffer = cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 70])
            if not ret:
                app.logger.error("Failed to encode frame")
                continue

            frame = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')
        frame_count += 1

    cap.release()

@app.route('/video_feed/<stream_id>')
def video_feed(stream_id):
    try:
        if stream_id == "1":
            def generate():
                global latest_frame
                while True:
                    if latest_frame is not None:
                        _, jpeg = cv2.imencode('.jpg', latest_frame)
                        frame = jpeg.tobytes()
                        yield (b'--frame\r\n'
                               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n\r\n')
                    else:
                        app.logger.warning("No latest frame available")
                        continue

            return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')
        elif stream_id == "2":
            return Response(generate_frames(STREAM_URL), mimetype='multipart/x-mixed-replace; boundary=frame')
        else:
            app.logger.error(f"Invalid stream ID: {stream_id}")
            return jsonify({'error': 'Invalid stream ID'}), 400
    except Exception as e:
        app.logger.error(f"Failed to stream video for ID {stream_id}: {str(e)}")
        return jsonify({'error': 'Failed to process video stream'}), 500

@app.route('/upload_video', methods=['POST'])
def upload_video():
    try:
        global latest_frame
        nparr = np.frombuffer(request.data, np.uint8)
        latest_frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        if latest_frame is None:
            raise ValueError("Failed to decode video frame")
        return 'Video uploaded successfully', 200
    except Exception as e:
        app.logger.error(f"Failed to upload video frame: {str(e)}")
        return 'Failed to upload video frame', 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True, threaded=True)
